from espdl import CocoDetector
from jpeg import Decoder, Encoder
from myufont import CustomBMFont
from machine import Pin,SDCard
import os
sd = SDCard(slot=0,width=4, sck=43, cmd=44, data=(39, 40, 41, 42))
os.mount(sd, '/sd')
decoder = Decoder()
encoder = Encoder(width=405, height=540,pixel_format="RGB888")
face_detector = CocoDetector(width=405, height=540)
MSCOCO_CLASSES = [
    "人", "自行车", "汽车", "摩托车", "飞机", "公共汽车", "火车", "卡车", "船", "交通灯",
    "消防栓", "消防水带", "停车计时器", "长椅", "鸟", "猫", "狗", "马", "羊", "牛",
    "大象", "熊", "斑马", "长颈鹿", "背包", "伞", "手提包", "领带", "行李箱", "飞盘",
    "滑雪板", "滑雪杖", "滑板", "冲浪板", "网球拍", "瓶子", "酒杯", "杯子", "刀叉", "碗",
    "香蕉", "苹果", "三明治", "橙子", "西兰花", "胡萝卜", "热狗", "披萨", "甜甜圈", "蛋糕",
    "椅子", "沙发", "盆栽", "床", "餐桌", "马桶", "电视", "笔记本电脑", "鼠标", "遥控器",
    "键盘", "手机", "微波炉", "烤箱", "烤面包机", "水槽", "冰箱", "书", "时钟", "花瓶",
    "剪刀", "泰迪熊", "吹风机", "牙刷"
]
font = CustomBMFont('/sd/text_full_16px_2312.v3.bmf')
# 捕获并处理图像
img = open("yolo.jpg", "rb").read()  # 捕获原始图像（通常是JPEG格式）
framebuffer = decoder.decode(img)  # 转换为RGB888格式
# 将memoryview转换为bytearray以便修改
framebuffer = bytearray(framebuffer)
# 运行人脸检测
results = face_detector.run(framebuffer)

# 绘制边框
def draw_rectangle(buffer, width, height, x, y, w, h,font,label, color=(255, 0, 0)):
    """
    在RGB888格式的图像缓冲区上绘制矩形边框
    :param buffer: 图像缓冲区
    :param width: 图像宽度
    :param height: 图像高度
    :param x: 边框左上角的x坐标
    :param y: 边框左上角的y坐标
    :param w: 边框宽度
    :param h: 边框高度
    :param color: 边框颜色（RGB格式）
    """
    # 辅助函数：设置单个像素的颜色
    def set_pixel(buffer, width, x, y, color):
        offset = (y * width + x) * 3
        buffer[offset] = color[0]  # R
        buffer[offset + 1] = color[1]  # G
        buffer[offset + 2] = color[2]  # B
    def is_chinese(ch):
        """判断一个字符是否为中文字符"""
        if '\u4e00' <= ch <= '\u9fff' or \
           '\u3400' <= ch <= '\u4dbf' or \
           '\u20000' <= ch <= '\u2a6df':
            return True
        return False
    def text(font, text, x_start, y_start, color,spacing=0, line_spacing=0, max_width=width):
        font_size = font.font_size
        bytes_per_row = (font_size + 7) // 8  # 每行占用的字节数
        x, y = x_start, y_start
        
        for char in text:
            # 处理换行符
            if char == '\n':
                y += font_size + line_spacing
                x = x_start
                continue
            if char == '\r':
                x += 2*font_size
                continue
            # 获取字符宽度(中文字符全宽，ASCII字符半宽)
            char_width = font_size if is_chinese(char) else font_size // 2
            
            # 检查是否需要换行
            if max_width is not None and x + char_width > x_start + max_width:
                y += font_size + line_spacing
                x = x_start
            
            # 获取字符位图
            bitmap = font.get_char_bitmap(char)
            
            # 绘制字符
            for row in range(font_size):
                for col in range(char_width if not is_chinese(char) else font_size):
                    byte_idx = row * bytes_per_row + col // 8
                    bit_mask = 0x80 >> (col % 8)
                    
                    if byte_idx < len(bitmap) and (bitmap[byte_idx] & bit_mask):
                        set_pixel(framebuffer,max_width,x + col, y + row, color)
            
            # 移动到下一个字符位置
            x += char_width + spacing
    # 绘制上边框
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y < height:
            set_pixel(buffer, width, i, y, color)

    # 绘制下边框
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y + h < height:
            set_pixel(buffer, width, i, y + h, color)

    # 绘制左边框
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x < width:
            set_pixel(buffer, width, x, j, color)

    # 绘制右边框
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x + w < width:
            set_pixel(buffer, width, x + w, j, color)
    text(font,label, x, y-20, color)

# 在图像上绘制人脸边框
for face in results:
    #print(face)
    x1, y1, x2, y2 = face['box']
    label = MSCOCO_CLASSES[face['category']]+":"+str(int(face['score']*100))+"%"
    draw_rectangle(framebuffer, 405, 540, x1, y1, x2-x1, y2-y1,font,label)  # 使用红色边框
    print(label)
# 将带有边框的图像重新编码为JPEG格式并保存
marked_img = encoder.encode(framebuffer)
with open("yolo_marked.jpg", "wb") as f:
    f.write(marked_img)
